#include "server.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <arpa/inet.h>
#include <sys/socket.h>
#include <netinet/in.h>

sqlite3 *db;

int init_server_socket() {
	    int server_socket = socket(AF_INET, SOCK_STREAM, 0);
	        if (server_socket == -1) {
			        perror("socket creation failed");
				        exit(EXIT_FAILURE);
					    }

		    struct sockaddr_in server_addr;
		        server_addr.sin_family = AF_INET;
			    server_addr.sin_addr.s_addr = INADDR_ANY;
			        server_addr.sin_port = htons(SERVER_PORT);

				    if (bind(server_socket, (struct sockaddr *)&server_addr, sizeof(server_addr)) {
						            perror("bind failed");
							            close(server_socket);
								            exit(EXIT_FAILURE);
									        }

										    if (listen(server_socket, 5)) {
										            perror("listen failed");
											            close(server_socket);
												            exit(EXIT_FAILURE);
													        }

														    return server_socket;
														    }
														    void log_event(sqlite3 *db, const char *event_type, const char *username, const char *details) {
														        char *err_msg = NULL;
															    char sql[512];
															        
															        if (username) {
																        snprintf(sql, sizeof(sql), 
																			                "INSERT INTO logs(event_type, username, details) VALUES('%s', '%s', '%s')",
																					                event_type, username, details);
																	    } else {
																		            snprintf(sql, sizeof(sql), 
																					                    "INSERT INTO logs(event_type, details) VALUES('%s', '%s')",
																							                    event_type, details);
																			        }
																    
																    if (sqlite3_exec(db, sql, NULL, 0, &err_msg) != SQLITE_OK) {
																	            fprintf(stderr, "SQL error: %s\n", err_msg);
																		            sqlite3_free(err_msg);
																			        }
														    }

void *handle_client(void *arg) {
	    ClientInfo *client = (ClientInfo *)arg;
	        int client_socket = client->client_socket;
		    
		    MessageHeader header;
		        unsigned char buffer[MAX_MSG_LEN + sizeof(MessageHeader)];
			    
			    while (1) {
				            
				            ssize_t bytes_read = read(client_socket, &header, sizeof(MessageHeader));
					            if (bytes_read <= 0) {
							                printf("Client disconnected\n");
									            log_event(db, "LOGOUT", client->username, "Client disconnected");
										                break;
												        }
						            
						            bytes_read = read(client_socket, buffer, header.msg_len);
							            if (bytes_read <= 0) {
									                printf("Error reading message body\n");
											            break;
												            }
								            
								            switch (header.msg_type) {
										                case MSG_TYPE_LOGIN: {
															                     LoginRequest *login_req = (LoginRequest *)buffer;
																	                     
																	                     char sql[256];
																			                     snprintf(sql, sizeof(sql), 
																							                             "SELECT password_hash, salt FROM users WHERE username='%s'", 
																										                             login_req->username);
																					                     
																					                     sqlite3_stmt *stmt;
																							                     if (sqlite3_prepare_v2(db, sql, -1, &stmt, NULL) != SQLITE_OK) {
																										                         printf("Database error\n");
																													                     break;
																															                     }
																									                     
																									                     LoginResponse resp;
																											                     if (sqlite3_step(stmt) == SQLITE_ROW) {
																														                         const unsigned char *db_hash = sqlite3_column_text(stmt, 0);
																																	                     const char *salt = (const char *)sqlite3_column_text(stmt, 1);
																																			                         
																																			                         if (memcmp(login_req->password_hash, db_hash, SM3_DIGEST_LENGTH) == 0) {
																																							                         resp.success = 1;
																																										                         strncpy(resp.salt, salt, SM3_DIGEST_LENGTH);
																																													                         strncpy(client->username, login_req->username, MAX_USERNAME_LEN);
																																																                         
																																																                         RAND_bytes(client->session_key, sizeof(client->session_key));
																																																			                         
																																																			                         log_event(db, "LOGIN", client->username, "Login successful");
																																																						                     } else {
																																																									                             resp.success = 0;
																																																												                             log_event(db, "LOGIN_FAIL", login_req->username, "Invalid password");
																																																															                         }
																																						                 } else {
																																									                     resp.success = 0;
																																											                         log_event(db, "LOGIN_FAIL", login_req->username, "User not found");
																																														                 }
																													                     
																													                     sqlite3_finalize(stmt);
																															                     
																															                     MessageHeader resp_header = {MSG_TYPE_LOGIN_RESP, sizeof(LoginResponse)};
																																	                     write(client_socket, &resp_header, sizeof(resp_header));
																																			                     write(client_socket, &resp, sizeof(resp));
																																					                     break;
																																							                 }
														                 case MSG_TYPE_CHAT: {
																			                     if (strlen(client->username) == 0) {
																						                         printf("Unauthorized chat attempt\n");
																									                     break;
																											                     }
																					                     
																					                     ChatMessage *chat_msg = (ChatMessage *)buffer;
																							                     unsigned char decrypted_msg[MAX_MSG_LEN];
																									                     
																									                     sm4_decrypt(client->session_key, chat_msg->encrypted_msg, 
																													                                 decrypted_msg, chat_msg->msg_len);
																											                     
																											                     printf("Message from %s to %s: %.*s\n", client->username, 
																															                            chat_msg->target_user, chat_msg->msg_len, decrypted_msg);
																													                     
																													                     log_event(db, "CHAT_SEND", client->username, chat_msg->target_user);
																															                     
																															                     
																															                     MessageHeader chat_header = {MSG_TYPE_CHAT, header.msg_len};
																																	                     write(client_socket, &chat_header, sizeof(chat_header));
																																			                     write(client_socket, buffer, header.msg_len);
																																					                     break;
																																							                 }
																		                 default:
																		                     printf("Unknown message type: %d\n", header.msg_type);
																				                     break;
																						             }
									        }
			        
			        close(client_socket);
				    free(client);
				        return NULL;
}

int main() {
	    
	    if (sqlite3_open("chat.db", &db) != SQLITE_OK) {
		            fprintf(stderr, "Can't open database: %s\n", sqlite3_errmsg(db));
			            return 1;
				        }
	        
	        int server_socket = init_server_socket();
		    printf("Server started on port %d\n", SERVER_PORT);
		        
		        while (1) {
				        struct sockaddr_in client_addr;
					        socklen_t addr_len = sizeof(client_addr);
						        
						        int client_socket = accept(server_socket, (struct sockaddr *)&client_addr, &addr_len);
							        if (client_socket < 0) {
									            perror("accept failed");
										                continue;
												        }
								        
								        printf("New connection from %s:%d\n", 
											               inet_ntoa(client_addr.sin_addr), ntohs(client_addr.sin_port));
									        
									        ClientInfo *client = malloc(sizeof(ClientInfo));
										        client->client_socket = client_socket;
											        memset(client->username, 0, MAX_USERNAME_LEN);
												        
												        pthread_t thread_id;
													        if (pthread_create(&thread_id, NULL, handle_client, (void *)client) != 0) {
															            perror("pthread_create failed");
																                close(client_socket);
																		            free(client);
																			            }
														        
														        pthread_detach(thread_id);
															    }
			    
			    sqlite3_close(db);
			        close(server_socket);
				    return 0;
}
